package com.mj.set.tree;

import java.util.Comparator;

public class AVLTree<E> extends BBST<E> {

    public AVLTree() {
        this(null);
    }

    /**
     * 调用父类的构造方法 传入comparator 比较器
     *
     * @param comparator
     */
    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    /**
     * 添加node之后的调整
     * 子类实现 不影响父类
     *
     * @param node 新添加的节点
     */
    @Override
    protected void afterAdd(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                //更新高度
                updateHeight(node);
            } else { // 找到第一个不平衡的节点
                //恢复平衡
                rebalance(node);
                break;
            }
        }
    }
    /**
     * 删除node之后的调整
     * 子类实现 不影响父类
     *
     * @param node 被删除的节点
     */
    @Override
    protected void afterRemove(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                //更新高度
                updateHeight(node);
            } else { // 找到第一个不平衡的节点
                //恢复平衡
                rebalance(node);
            }
        }
    }

    /**
     * 实现返回自己的节点
     *
     * @param element
     * @param parent
     * @return
     */
    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element, parent);
    }

    /**
     * 恢复平衡
     *
     * @param grand 高度最低的不平衡节点
     */
    private void rebalance2(Node<E> grand) {
        Node<E> parent = ((AVLNode<E>) grand).tallerChild();
        Node<E> node = ((AVLNode<E>) parent).tallerChild();
        if (parent.isLeftChild()) {//L
            if (node.isLeftChild()) {//LL
                rotateRight(grand);
            } else {//LR
                rotateLeft(parent);
                rotateRight(grand);
            }
        } else {//R
            if (node.isLeftChild()) {//RL
                rotateRight(parent);
                rotateLeft(grand);
            } else {//RR
                rotateLeft(grand);
            }
        }
    }

    /**
     * 恢复平衡
     * @param grand
     */
    private void rebalance(Node<E> grand) {
        Node<E> parent = ((AVLNode<E>) grand).tallerChild();
        Node<E> node = ((AVLNode<E>) parent).tallerChild();
        if (parent.isLeftChild()) {//L
            if (node.isLeftChild()) {//LL
                rotate(grand,node,node.right,parent, parent.right, grand );
            } else {//LR
                rotate(grand,parent,node.left, node, node.right , grand);
            }
        } else {//R
            if (node.isLeftChild()) {//RL
                rotate(grand,grand,node.left,node, node.right, parent );
            } else {//RR
                rotate(grand,grand,parent.left,parent, node.left, node);
            }
        }
    }

    @Override
    protected void rotate(Node<E> r, Node<E> b, Node<E> c, Node<E> d, Node<E> e, Node<E> f) {
        super.rotate(r, b, c, d, e, f);
        //左右子树发生了变化，应该更新父节点的高度
        updateHeight(b);
        updateHeight(f);
        updateHeight(d);
    }

    /**
     * 更新p、T2、g的 parent 属性
     *
     *公共代码:不管是左旋转、右旋转，都要执行的
     *@param grand 失衡节点
     *@param parent 失衡节点的tallerChild
     *@param child g和p需要交换的子树（本来是p的子树，后面会变成g的子树)
     */
    @Override
    protected void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {

      super.afterRotate(grand, parent, child);
        //更新高度
        updateHeight(grand);
        updateHeight(parent);
    }

    /**
     * 每个节点的平衡因子只可能是 `1、0、-1`（绝对值 ≤ 1，如果超过 1，称之为“失衡”）
     *
     * @param node
     * @return
     */
    private boolean isBalanced(Node<E> node) {
        return Math.abs(((AVLNode<E>) node).balanceFactor()) <= 1;
    }

    private void updateHeight(Node<E> node) {
        ((AVLNode<E>) node).updateHeight();
    }

    private static class AVLNode<E> extends Node<E> {
        int height = 1;//树的高度

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        /**
         * 平衡因子（Balance Factor）：某结点的左右子树的高度差 (左子树的高度 -  右子树的高度)
         *
         * @return
         */
        public int balanceFactor() {
            //强制将左子节点转换为AVL节点
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            return leftHeight - rightHeight;
        }

        public void updateHeight() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        /**
         * 获取高度最高的子节点
         * @return
         */
        public Node<E> tallerChild() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            if (leftHeight > rightHeight) return left;
            if (leftHeight < rightHeight) return right;
            //如果 高度相等 返回父节点方向的节点
            return isLeftChild() ? left : right;
        }

        @Override
        public String toString() {
            String parentString = "null";
            if (parent != null) {
                parentString = parent.element.toString();
            }
            return element + "_p(" + parentString + ")_h("+height+")";
        }
    }
}
